from django.core.management.base import BaseCommand
from django.db import connection
from parts.models import Parts
from datetime import datetime, timedelta


class Command(BaseCommand):
    help = 'Initialize the parts database with partitions'
    
    def handle(self, *args, **kwargs):
        try:
            with connection.cursor() as cursor:
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS parts (
                        id SERIAL,
                        sn VARCHAR(32) NOT NULL,
                        create_time TIMESTAMPTZ NOT NULL,
                        status SMALLINT,
                        PRIMARY KEY (id, create_time)
                    ) PARTITION BY RANGE (create_time);
                """)
            
            # 创建未来3个月的分区
            now = datetime.now()
            for month_offset in range(3):
                dt = now + timedelta(days=30*month_offset)
                Parts.create_partition(dt.year, dt.month)
                self.stdout.write(self.style.SUCCESS(f'Created partition for {dt.year}-{dt.month:02d}'))
            
            self.stdout.write(self.style.SUCCESS('Database initialized successfully'))
        
        except Exception as e:
            self.stdout.write(self.style.ERROR(f'Error: {str(e)}'))